package com.yupi.loveplanneraiagent.tools;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpUtil;
import com.yupi.loveplanneraiagent.constant.FileConstant;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.io.File;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.regex.Pattern;

/**
 * 资源下载工具
 */
public class ResourceDownloadTool {

    // 定义允许的协议
    private static final String[] ALLOWED_PROTOCOLS = {"http", "https"};

    // 定义安全的文件扩展名白名单
    private static final String[] ALLOWED_EXTENSIONS = {
            ".txt", ".pdf", ".jpg", ".jpeg", ".png", ".gif", ".bmp",
            ".doc", ".docx", ".xls", ".xlsx", ".ppt", ".pptx",
            ".zip", ".rar", ".tar", ".gz", ".mp3", ".mp4", ".avi"
    };

    @Tool(description = "Download a resource from a given URL")
    public String downloadResource(
            @ToolParam(description = "URL of the resource to download") String url,
            @ToolParam(description = "Name of the file to save the downloaded resource") String fileName) {

        try {
            // 参数验证
            if (StrUtil.isBlank(url) || StrUtil.isBlank(fileName)) {
                return "Error: Both URL and file name are required.";
            }

            // URL安全性验证
            if (!isValidUrl(url)) {
                return "Error: Invalid or unsafe URL provided.";
            }

            // 文件名安全性验证
            if (!isValidFileName(fileName)) {
                return "Error: Invalid file name provided.";
            }

            // 验证文件扩展名
            if (!isValidFileExtension(fileName)) {
                return "Error: File type not allowed. Only certain file types are supported for download.";
            }

            String fileDir = FileConstant.FILE_SAVE_DIR + "/download";
            String filePath = fileDir + "/" + fileName;

            // 创建目录
            FileUtil.mkdir(fileDir);

            // 验证目标路径安全性（防止路径遍历）
            File targetFile = new File(filePath);
            String canonicalPath = targetFile.getCanonicalPath();
            String canonicalDir = new File(fileDir).getCanonicalPath();

            if (!canonicalPath.startsWith(canonicalDir)) {
                return "Error: Unsafe file path detected.";
            }

            // 使用 Hutool 的 downloadFile 方法下载资源
            HttpUtil.downloadFile(url, targetFile);

            return "Resource downloaded successfully to: " + filePath;

        } catch (Exception e) {
            return "Error downloading resource: " + e.getMessage();
        }
    }

    /**
     * 验证URL的安全性
     */
    private boolean isValidUrl(String url) {
        try {
            URL parsedUrl = new URL(url);

            // 检查协议
            String protocol = parsedUrl.getProtocol();
            boolean protocolAllowed = false;
            for (String allowedProtocol : ALLOWED_PROTOCOLS) {
                if (allowedProtocol.equals(protocol)) {
                    protocolAllowed = true;
                    break;
                }
            }

            if (!protocolAllowed) {
                return false;
            }

            // 阻止本地地址访问（防止SSRF）
            String host = parsedUrl.getHost();
            if (isLocalAddress(host)) {
                return false;
            }

            return true;
        } catch (MalformedURLException e) {
            return false;
        }
    }

    /**
     * 检查是否为本地地址（防止SSRF攻击）
     */
    private boolean isLocalAddress(String host) {
        // 检查IP地址是否为内网地址
        if (host.matches("^127\\.\\d+\\.\\d+\\.\\d+$") ||  // 127.x.x.x
                host.matches("^10\\.\\d+\\.\\d+\\.\\d+$") ||   // 10.x.x.x
                host.matches("^172\\.(1[6-9]|2[0-9]|3[01])\\.\\d+\\.\\d+$") || // 172.16-31.x.x
                host.matches("^192\\.168\\.\\d+\\.\\d+$") ||  // 192.168.x.x
                "localhost".equalsIgnoreCase(host) ||
                host.startsWith("internal.") ||  // 假设internal.开头的为内网地址
                host.startsWith("intranet.")) {  // 假设intranet.开头的为内网地址
            return true;
        }
        return false;
    }

    /**
     * 验证文件名的安全性
     */
    private boolean isValidFileName(String fileName) {
        // 检查是否包含路径遍历字符
        if (fileName.contains("..") || fileName.contains("../") ||
                fileName.contains("..\\") || fileName.contains("./") || fileName.contains(".\\")) {
            return false;
        }

        // 检查是否包含危险字符
        Pattern dangerousChars = Pattern.compile("[<>:\"/\\\\|?*]");
        if (dangerousChars.matcher(fileName).find()) {
            return false;
        }

        return true;
    }

    /**
     * 验证文件扩展名是否在白名单中
     */
    private boolean isValidFileExtension(String fileName) {
        String lowerCaseFileName = fileName.toLowerCase();
        for (String ext : ALLOWED_EXTENSIONS) {
            if (lowerCaseFileName.endsWith(ext)) {
                return true;
            }
        }
        return false;
    }
}



